import base64
import hashlib
import json
import os
from urllib.parse import urlparse, parse_qs

import requests
from DrissionPage import ChromiumPage, ChromiumOptions

# DefaultUserAgent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/127.0.0.0 Safari/537.36"
DefaultUserAgent = "ChatGPT/1.2024.206 (Mac OS X 14.4.1; arm64; build 1722395003)"


def generate_code_verifier():
    token = os.urandom(32)
    code_verifier = base64.urlsafe_b64encode(token).rstrip(b'=')
    return code_verifier.decode('utf-8')


def generate_code_challenge(code_verifier):
    m = hashlib.sha256()
    m.update(code_verifier.encode('utf-8'))
    code_challenge = base64.urlsafe_b64encode(m.digest()).rstrip(b'=')
    return code_challenge.decode('utf-8')


class AuthBrowserios():

    def __init__(self):

        browser_path = "/usr/bin/google-chrome"
        options = ChromiumOptions()
        options.set_paths(browser_path=browser_path)

        if DefaultUserAgent:
            options.set_user_agent(DefaultUserAgent)

        arguments = []

        for argument in arguments:
            options.set_argument(argument)

        options.headless(False)

        self.driver = ChromiumPage(addr_or_opts=options)
        self.driver.set.cookies.clear()

    def get_preauth_cookie(self):
        url = "https://publicapi.closeai.biz/auth/preauth"
        res = requests.get(url).json()
        return res["preauth_cookie"]

    def login(self):

        code_verifier = generate_code_verifier()
        code_challenge = generate_code_challenge(code_verifier)

        base_url = "https://auth0.openai.com/authorize?"
        client_info = {"version": "2.9.0-beta.0", "name": "Auth0.swift", "env": {"swift": "5.x", "macOS": "14.4"}}

        params = {
            "response_type": "code",
            "client_id": "pdlLIX2Y72MIl2rhLhTE9VV9bN905kBh",
            "code_challenge": code_challenge,
            "code_challenge_method": "S256",
            "scope": "openid%20email%20profile%20offline_access%20model.request%20model.read%20organization.read%20organization.write",
            "redirect_uri": "com.openai.chat://auth0.openai.com/ios/com.openai.chat/callback",
            "audience": "https://api.openai.com/v1",
            "preauth_cookie": self.get_preauth_cookie(),
            "ios_app_version": "1723258889",
            "prompt": "login",
            "auth0Client": base64.b64encode(json.dumps(client_info).replace(" ", "").encode()).decode()
        }

        query_string = "&".join([f"{key}={value}" for key, value in params.items()])
        url = base_url + query_string
        # print(url)

        tab_id = self.driver.new_tab(url).tab_id
        tab = self.driver.get_tab(tab_id)
        tab.listen.start('com.openai.chat://auth0.openai.com/ios/com.openai.chat/callback')

        res = tab.listen.wait()
        parsed_url = urlparse(res.url)

        query_params = parse_qs(parsed_url.query)
        code = query_params.get('code', [None])[0]

        print("code", code)
        print("code_verifier", code_verifier)
        # self.get_token(code, code_verifier)

    def get_token(self, code, code_verifier):
        url = 'https://auth0.openai.com/oauth/token'
        data = {
            'redirect_uri': 'com.openai.chat://auth0.openai.com/ios/com.openai.chat/callback',
            'grant_type': 'authorization_code',
            'client_id': 'pdlLIX2Y72MIl2rhLhTE9VV9bN905kBh',
            'code': code,
            'code_verifier': code_verifier,
        }
        resp = requests.post(url, json=data, allow_redirects=False)
        print(resp.json())
